# -*- coding: utf-8 -*-
"""
Created on Mon Mar 22 13:44:05 2021

@author: ckpark
"""
import os
os.environ["PROJ_LIB"] = r'C:\Users\ckpark\anaconda3\Library\share'

from netCDF4 import Dataset
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import sys
import matplotlib.gridspec as gridspec
from scipy.stats import linregress
import matplotlib.colors as mc
from matplotlib.patches import Polygon
import matplotlib as m
import matplotlib
from scipy import stats
import scipy.signal
import scipy.signal as signal

#################################################################
def moving(inte, n):
    result = np.convolve(inte, np.ones((n,))/n, mode='same')
    return result
def init_plotting():
    plt.rcParams['font.family'] = 'Times New Roman'
init_plotting()
######################################

ceyear = 2021
csyear = ceyear - 29

read = np.loadtxt('../DATA/BIN/GLDAS_47156', dtype='str')
gyr = np.array(read[:,0], dtype='int')
gmo = np.array(read[:,1], dtype='int')
gldas = np.array(read[:,2], dtype='float')

read = np.loadtxt('./RESULT/EDI_47156', dtype='str')
yr = np.array(read[1:,0], dtype='int')
mo = np.array(read[1:,1], dtype='int')
dy = np.array(read[1:,2], dtype='int')
edi = np.array(read[1:,3], dtype='float')
pre = np.array(read[1:,4], dtype='float')

syear = yr[0]
eyear = yr[-1]
nyear = eyear - syear + 1
nt = nyear * 12

me_dat = np.zeros((nyear,12,31))
mp_dat = np.zeros((nyear,12,31))
me_dat[:,:,:] = -999.0
mp_dat[:,:,:] = -999.0
for t in range(len(yr)):
    me_dat[yr[t]-syear,mo[t]-1,dy[t]-1] = edi[t]
    mp_dat[yr[t]-syear,mo[t]-1,dy[t]-1] = pre[t]
    
m_edi = []
m_pre = []
myr = []
for y in range(nyear):
    for m in range(12):
        edat = me_dat[y,m,:]
        pdat = mp_dat[y,m,:]
        
        edat = np.ma.masked_values(edat, -999.0)
        pdat = np.ma.masked_values(pdat, -999.0)
        
        m_edi.append(np.min(edat)); m_pre.append(np.sum(pdat))
        myr.append(yr[0]+y)
########################
d_pre = []
for t in range(len(m_pre)):
    if myr[t] >= csyear and myr[t] <= ceyear:
        d_pre.append(m_pre[t])
        
c_pre = []
for t in range(12):
    c_pre.append(np.mean(d_pre[t::12]))

a_pre = []; imo = 0
for t in range(len(m_pre)):
    adat = m_pre[t] - c_pre[imo]
    a_pre.append(adat)
    
    imo += 1
    if imo == 12:
        imo = 0

r_pre = moving(a_pre,12)
#########################
d_gld = []
for t in range(len(gyr)):
    if gyr[t] >= csyear and gyr[t] <= ceyear:
        d_gld.append(gldas[t])

for t in range(len(gyr)):
    if gyr[t] == syear:
        st = t
        break
gldas = gldas[st:]
    
c_gld = []
for t in range(12):
    c_gld.append(np.mean(d_gld[t::12]))

a_gld = []; imo = 0
for t in range(len(gldas)):
    adat = gldas[t] - c_gld[imo]
    a_gld.append(adat)
    
    imo += 1
    if imo == 12:
        imo = 0

r_gld = moving(a_gld,12)  
a_gld.append(-999.0)
a_gld = np.ma.masked_values(a_gld, -999.0)
   
 
############## FIG #############
fig = plt.figure(figsize=(10,10))
for xx in range(1):
    ax1 = fig.add_subplot(3,1,xx+1)

    xax = np.arange(0,nt)
    ax1.set_xlim(0,nt-1)

    xlab = np.arange(syear,eyear+1,5)
    xtick = np.arange(0,nt,12*5)

    ax1.set_xticks(xtick)
    ax1.set_xticklabels(xlab)
    ax1.set_xlabel('Year',fontsize=26)    
    #plt.xticks(rotation=45)

  
    ax1.set_yticks(np.arange(-10.0,10.0,1.0))
    ax1.set_ylim(-3.5,3.4)
        
    ax1.axhline(y=0.0, color='black', linestyle='solid')
    ax1.axhline(y=-2.0, color='red', linestyle='dashed')

    ax1.plot(xax,m_edi,color='black',zorder=1,alpha=1.0)
    #ax1.plot(xax,macd_rpdo,color='blue',zorder=1,alpha=1.0)
    ax1.set_ylabel('Index',fontsize=24)
    ax1.tick_params(labelsize=23)
        
plt.subplots_adjust(hspace=0.0,wspace=0.2)
plt.show()
outfig = '../99.FIG/EDI_time_series.png'      
fig.savefig(outfig,bbox_inches='tight') 
#############################################################33
############## FIG #############
fig = plt.figure(figsize=(10,10))
for xx in range(2):
    ax1 = fig.add_subplot(3,1,xx+1)

    xax = np.arange(0,nt)
    ax1.set_xlim(0,nt-1)
    
    if xx < 1:
        xlab = np.arange(0,1)
        xtick = np.arange(1000,1000+nt,100*5)
    if xx == 1:
        xlab = np.arange(syear,eyear+1,5)
        xtick = np.arange(0,nt,12*5)

    ax1.set_xticks(xtick)
    ax1.set_xticklabels(xlab)
    ax1.set_xlabel('Year',fontsize=26)    
    #plt.xticks(rotation=45)

    if xx == 2:    
        ax1.set_yticks(np.arange(-10.0,10.0,1.0))
        ax1.set_ylim(-3.5,3.4)
        
        ax1.axhline(y=0.0, color='black', linestyle='solid')
        ax1.axhline(y=-2.0, color='red', linestyle='dashed')

        ax1.plot(xax,m_edi,color='black',zorder=1,alpha=1.0)
        #ax1.plot(xax,macd_rpdo,color='blue',zorder=1,alpha=1.0)
        ax1.set_ylabel('Index',fontsize=24)
        ax1.tick_params(labelsize=23)
        
    if xx == 0:
        ax1.set_yticks(np.arange(-225,300,75))
        ax1.set_ylim(-220,220)
        
        ax1.axhline(y=0.0, color='black', linestyle='solid')
        #ax1.axhline(y=-2.0, color='red', linestyle='dashed')

        for t in range(len(xax)):
            if a_pre[t] > 0:
                ax1.bar(xax[t],a_pre[t],color='blue',zorder=1,alpha=1.0)
            else:
                ax1.bar(xax[t],a_pre[t],color='red',zorder=1,alpha=1.0)
        ax1.plot(xax,r_pre,color='black',zorder=1,alpha=1.0)
        #ax1.plot(xax,macd_rpdo,color='blue',zorder=1,alpha=1.0)
        ax1.set_ylabel('Precipitation'+'\n'+'(mm)',fontsize=24)
        ax1.tick_params(labelsize=23)

    if xx == 1:
        ax1.set_yticks(np.arange(-300,300,50))
        ax1.set_ylim(-120,140)
        
        ax1.axhline(y=0.0, color='black', linestyle='solid')
        #ax1.axhline(y=-2.0, color='red', linestyle='dashed')

        for t in range(len(xax)):
            if a_gld[t] > 0:
                ax1.bar(xax[t],a_gld[t],color='blue',zorder=1,alpha=1.0)
            else:
                ax1.bar(xax[t],a_gld[t],color='red',zorder=1,alpha=1.0)
        ax1.plot(xax[:-1],r_gld[:],color='black',zorder=1,alpha=1.0)
        #ax1.plot(xax,macd_rpdo,color='blue',zorder=1,alpha=1.0)
        ax1.set_ylabel('Soil moisutre'+'\n'+'(mm)',fontsize=24)
        ax1.tick_params(labelsize=23)

plt.subplots_adjust(hspace=0.0,wspace=0.2)
plt.show()
outfig = '../99.FIG/PRE_SOIL_time_series.png'      
fig.savefig(outfig,bbox_inches='tight') 
#############################################################33
